#include "../include/KeyFrameDatabase.hpp"
namespace ORBSLAM
{
    KeyFrameDatabase::KeyFrameDatabase(/* args */)
    {
    }

    KeyFrameDatabase::~KeyFrameDatabase()
    {
    }

    KeyFrameDatabase::KeyFrameDatabase(const DBoW3::Vocabulary &vocabulary)
    {
        mp_vocabulary = &vocabulary;
        cout << "vocabulary size:" << mp_vocabulary->size() << endl;
        mvlp_invert_index.resize(mp_vocabulary->size());
    }

    void KeyFrameDatabase::Add(KeyFrame *p_keyframe)
    {
        unique_lock<mutex> lock(mmutex_keyframe_db);
        for (auto bow : p_keyframe->mv_bowvector)
        {
            mvlp_invert_index[bow.first].push_back(p_keyframe);
        }
    }

    void KeyFrameDatabase::Erase(KeyFrame *p_keyframe)
    {
        unique_lock<mutex> lock(mmutex_keyframe_db);
        for (auto bow : p_keyframe->mv_bowvector)
        {
            list<KeyFrame *> &lp_keyframe = mvlp_invert_index[bow.first];
            for (auto lit = lp_keyframe.begin(); lit != lp_keyframe.end(); ++lit)
            {
                if (p_keyframe == *lit)
                {
                    lp_keyframe.erase(lit);
                    break;
                }
            }
        }
    }

    bool Comp_first(const pair<float, KeyFrame *> &a, const pair<float, KeyFrame *> &b)
    {
        return a.first > b.first;
    }

    bool Comp_id(KeyFrame *p_keyframe1, KeyFrame *p_keyframe2)
    {
        return p_keyframe1->mn_keyframe_id < p_keyframe2->mn_keyframe_id;
    }

    void KeyFrameDatabase::Detect_nbest_candidates(KeyFrame *p_keyframe, vector<KeyFrame *> &vp_loop_cand, vector<KeyFrame *> &vp_merge_cand, int candidates_num)
    {
        list<KeyFrame *> lp_share_worlds;
        set<KeyFrame *> sp_connect_keyframe;
        {
            unique_lock<mutex> lock(mmutex_keyframe_db);
            sp_connect_keyframe = p_keyframe->Get_connect_keyframe();
            for (auto bow : p_keyframe->mv_bowvector)
            {
                list<KeyFrame *> lp_keyframe = mvlp_invert_index[bow.first];
                for (auto p_keyframe1 : lp_keyframe)
                {
                    if (p_keyframe1->mn_place_recognition_query != p_keyframe->mn_keyframe_id)
                    {
                        p_keyframe1->mn_place_recognition_worlds = 0;
                        // if (!sp_connect_keyframe.count(p_keyframe1))
                        // {
                        //     p_keyframe1->mn_place_recognition_query = p_keyframe->mn_keyframe_id;
                        //     lp_share_worlds.push_back(p_keyframe1);
                        // }
                        p_keyframe1->mn_place_recognition_query = p_keyframe->mn_keyframe_id;
                        lp_share_worlds.push_back(p_keyframe1);
                    }
                    p_keyframe1->mn_place_recognition_worlds++;
                }
            }
        }

        if (lp_share_worlds.empty())
        {
            return;
        }

        int max_common_worlds = 0;
        KeyFrame *p_keyframe_share_most_worlds;
        for (auto p_keyframe1 : lp_share_worlds)
        {
            if (p_keyframe1->mn_feature_num > 2000)
            {
                if (p_keyframe1->mn_place_recognition_worlds / 5 > max_common_worlds)
                {
                    max_common_worlds = p_keyframe1->mn_place_recognition_worlds / 5;
                    p_keyframe_share_most_worlds = p_keyframe1;
                }
            }
            else
            {
                if (p_keyframe1->mn_place_recognition_worlds > max_common_worlds)
                {
                    max_common_worlds = p_keyframe1->mn_place_recognition_worlds;
                    p_keyframe_share_most_worlds = p_keyframe1;
                }
            }
        }



        if (p_keyframe->mn_frame_id > 3500)
        {
            lp_share_worlds.sort(Comp_id);
            for (auto p_keyframei : lp_share_worlds)
            {
                cout << p_keyframei->mn_keyframe_id << ":" << p_keyframei->mn_place_recognition_worlds << endl;
            }
        }

        int min_common_worlds = max_common_worlds * 0.5f;
        list<pair<float, KeyFrame *>> lpair_score_keyframe;

        int n_score = 0;
        for (auto p_keyframe1 : lp_share_worlds)
        {
            if (p_keyframe1->mn_feature_num > 2000)
            {
                if (p_keyframe1->mn_place_recognition_worlds >= 5 * min_common_worlds)
                {
                    n_score++;
                    float si = mp_vocabulary->score(p_keyframe1->mv_bowvector, p_keyframe->mv_bowvector);
                    p_keyframe1->mf_place_recognition_score = si;
                    lpair_score_keyframe.push_back(make_pair(si, p_keyframe1));
                }
            }
            else
            {
                if (p_keyframe1->mn_place_recognition_worlds >= min_common_worlds)
                {
                    n_score++;
                    float si = mp_vocabulary->score(p_keyframe1->mv_bowvector, p_keyframe->mv_bowvector);
                    p_keyframe1->mf_place_recognition_score = si;
                    lpair_score_keyframe.push_back(make_pair(si, p_keyframe1));
                }
            }
        }

        if (lpair_score_keyframe.empty())
        {
            return;
        }

        list<pair<float, KeyFrame *>> lpair_accscore_keyframe;
        float best_accscore = 0;
        set<KeyFrame *> sp_tmp_keyframe;

        for (auto pair_score_keyframe : lpair_score_keyframe)
        {
            KeyFrame *p_keyframe1 = pair_score_keyframe.second;
            vector<KeyFrame *> vp_keyframe1_neighs = p_keyframe1->Get_best_covisibility_keyframe(10);
            float best_score = pair_score_keyframe.first;
            float accscore = best_score;
            KeyFrame *p_best_keyframe = p_keyframe1;
            bool b_abort = false;
            for (auto p_keyframe2 : vp_keyframe1_neighs)
            {
                if (sp_connect_keyframe.count(p_keyframe2))
                {
                    b_abort = true;
                    break;
                }
                if (p_keyframe2->mn_place_recognition_query != p_keyframe->mn_keyframe_id || p_keyframe2->mn_place_recognition_worlds < min_common_worlds)
                {
                    continue;
                }
                p_keyframe2->mf_place_recognition_score = mp_vocabulary->score(p_keyframe2->mv_bowvector, p_keyframe->mv_bowvector);
                accscore += p_keyframe2->mf_place_recognition_score;
                if (p_keyframe2->mf_place_recognition_score > best_score)
                {
                    best_score = p_keyframe2->mf_place_recognition_score;
                    p_best_keyframe = p_keyframe2;
                }
            }
            if (b_abort)
            {
                continue;
            }
            for (auto p_keyframek : p_best_keyframe->Get_best_covisibility_keyframe(5))
            {
                if (sp_connect_keyframe.count(p_keyframek))
                {
                    b_abort = true;
                    break;
                }
            }
            if (b_abort)
            {
                continue;
            }
            if (!sp_tmp_keyframe.count(p_best_keyframe) && !sp_connect_keyframe.count(p_best_keyframe))
            {
                sp_tmp_keyframe.insert(p_best_keyframe);
                lpair_accscore_keyframe.push_back(make_pair(accscore, p_best_keyframe));
            }
            if (accscore > best_accscore)
            {
                best_accscore = accscore;
            }
        }

        if (lpair_accscore_keyframe.empty())
        {
            return;
        }

        lp_share_worlds.sort(Comp_id);

        for (auto p_keyframei : lp_share_worlds)
        {

            cout << p_keyframei->mn_keyframe_id << ":" << p_keyframei->mn_place_recognition_worlds << endl;
        }

        lpair_accscore_keyframe.sort(Comp_first);
        vp_loop_cand.reserve(candidates_num);
        vp_merge_cand.reserve(candidates_num);
        set<KeyFrame *> sp_already_added_keyframe;
        int i = 0;
        list<pair<float, KeyFrame *>>::iterator it = lpair_accscore_keyframe.begin();
        while (i < lpair_accscore_keyframe.size() && (vp_loop_cand.size() < candidates_num || vp_merge_cand.size() < candidates_num))
        {
            KeyFrame *p_keyframe1 = it->second;
            if (p_keyframe1->Is_bad())
            {
                continue;
            }
            if (!sp_already_added_keyframe.count(p_keyframe1))
            {
                if (p_keyframe->Get_map() == p_keyframe1->Get_map() && vp_loop_cand.size() < candidates_num)
                {
                    vp_loop_cand.push_back(p_keyframe1);
                }
                else if (p_keyframe1->Get_map() != p_keyframe->Get_map() && vp_merge_cand.size() < candidates_num && !p_keyframe1->Get_map()->Is_bad())
                {
                    vp_merge_cand.push_back(p_keyframe1);
                }
                sp_already_added_keyframe.insert(p_keyframe1);
            }
            i++;
            it++;
        }
    }

    vector<KeyFrame *> KeyFrameDatabase::Detect_relocalization_candidates(Frame *p_frame, Map *p_map)
    {
        list<KeyFrame *> lp_keyframe_share_world;
        {
            unique_lock<mutex> lock(mmutex_keyframe_db);
            for (auto bow : p_frame->mv_bowvector)
            {
                list<KeyFrame *> &lp_keyframe = mvlp_invert_index[bow.first];
                for (auto p_keyframe1 : lp_keyframe)
                {
                    if (p_keyframe1->mn_relocal_query != p_frame->mn_frame_id)
                    {
                        p_keyframe1->mn_relocal_query = p_frame->mn_frame_id;
                        p_keyframe1->mn_relocal_worlds = 0;
                        lp_keyframe_share_world.push_back(p_keyframe1);
                    }
                    p_keyframe1->mn_relocal_worlds++;
                }
            }
        }
        if (lp_keyframe_share_world.empty())
        {
            return vector<KeyFrame *>();
        }

        int max_common_worlds = 0;
        for (auto p_keyframe1 : lp_keyframe_share_world)
        {
            if (p_keyframe1->mn_relocal_worlds > max_common_worlds)
            {
                max_common_worlds = p_keyframe1->mn_relocal_worlds;
            }
        }

        int min_common_worlds = 0.8f * max_common_worlds;

        list<pair<float, KeyFrame *>> lpair_score_and_match;
        int nscore = 0;

        for (auto p_keyframe1 : lp_keyframe_share_world)
        {

            if (p_keyframe1->mn_relocal_worlds > min_common_worlds)
            {
                nscore++;
                float si = mp_vocabulary->score(p_frame->mv_bowvector, p_keyframe1->mv_bowvector);
                p_keyframe1->mf_relocal_score = si;
                lpair_score_and_match.push_back(make_pair(si, p_keyframe1));
            }
        }

        if (lpair_score_and_match.empty())
        {
            return vector<KeyFrame *>();
        }

        list<pair<float, KeyFrame *>> lpair_accscore_and_match;
        float best_accscore = 0;

        for (auto pair_score_keyframe : lpair_score_and_match)
        {
            KeyFrame *p_keyframe1 = pair_score_keyframe.second;
            vector<KeyFrame *> vp_neigh_keyframe = p_keyframe1->Get_best_covisibility_keyframe(10);

            float best_score = pair_score_keyframe.first;
            float accscore = best_score;
            KeyFrame *p_best_keyframe = p_keyframe1;
            for (auto p_keyframe2 : vp_neigh_keyframe)
            {
                if (p_keyframe2->mn_relocal_query != p_frame->mn_frame_id)
                {
                    continue;
                }
                p_keyframe2->mf_relocal_score = mp_vocabulary->score(p_frame->mv_bowvector, p_keyframe2->mv_bowvector);
                accscore += p_keyframe2->mf_relocal_score;
                if (p_keyframe2->mf_relocal_score > best_score)
                {
                    p_best_keyframe = p_keyframe2;
                    best_score = p_keyframe2->mf_relocal_score;
                }
            }
            lpair_accscore_and_match.push_back(make_pair(accscore, p_best_keyframe));
            if (accscore > best_accscore)
            {
                best_accscore = accscore;
            }
        }

        float min_score_retain = 0.75f * best_accscore;
        set<KeyFrame *> sp_already_added_keyframe;
        vector<KeyFrame *> vp_relocal_candidates;
        vp_relocal_candidates.reserve(lpair_accscore_and_match.size());
        for (auto pair_accscore_keyframe : lpair_accscore_and_match)
        {
            const float &si = pair_accscore_keyframe.first;
            if (si > min_score_retain)
            {
                KeyFrame *p_keyframe1 = pair_accscore_keyframe.second;
                if (p_keyframe1->Get_map() != p_map)
                {
                    continue;
                }
                if (!sp_already_added_keyframe.count(p_keyframe1))
                {
                    vp_relocal_candidates.push_back(p_keyframe1);
                    sp_already_added_keyframe.insert(p_keyframe1);
                }
            }
        }

        return vp_relocal_candidates;
    }
}